package org.ruoyi.common.chat.xunfei;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.web.socket.WebSocketSession;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * 讯飞AI WebSocket客户端
 * 用于建立与讯飞AI的WebSocket连接并进行通信
 *
 * @author ruoyi
 */
@Slf4j
public class XunfeiWebSocketClient {

    private final String appId;
    private final String apiKey;
    private final String apiSecret;
    private final String hostUrl;
    private final WebSocketSession webSocketSession;
    private final List<RoleContent> historyList = new ArrayList<>();
    private WebSocket webSocket;
    private boolean isConnected = false;

    public XunfeiWebSocketClient(String appId, String apiKey, String apiSecret, String hostUrl, WebSocketSession webSocketSession) {
        this.appId = appId;
        this.apiKey = apiKey;
        this.apiSecret = apiSecret;
        this.hostUrl = hostUrl;
        this.webSocketSession = webSocketSession;
    }

    /**
     * 连接到讯飞AI
     *
     * @param url 连接地址
     */
    public void connect(String url) {
        // 添加URL参数检查
        if (StrUtil.isBlank(url)) {
            log.error("连接URL不能为空");
            return;
        }

        if (isConnected) {
            return;
        }

        OkHttpClient client = new OkHttpClient.Builder()
                .pingInterval(20, TimeUnit.SECONDS)  // 添加心跳机制
                .build();
        Request request = new Request.Builder()
                .url(url)
                .build();
        webSocket = client.newWebSocket(request, new XfWebSocketListener());
        isConnected = true;
    }

    /**
     * 发送消息到讯飞AI
     *
     * @param text 用户输入文本
     */
    public void sendTextToXunfei(String text) {
        sendTextToXunfei(text, null);
    }
    
    /**
     * 发送消息到讯飞AI（支持上下文）
     *
     * @param text 用户输入文本
     * @param context 上下文信息
     */
    public void sendTextToXunfei(String text, JSONArray context) {
        try {
            // 如果提供了上下文，则更新历史记录
            if (context != null && !context.isEmpty()) {
                synchronized (historyList) {
                    historyList.clear(); // 清除现有历史记录
                    // 根据上下文重建历史记录
                    for (int i = 0; i < context.size(); i++) {
                        JSONObject msg = context.get(i, JSONObject.class);
                        String role = msg.getStr("role");
                        String content = msg.getStr("content");
                        if (StrUtil.isNotBlank(role) && StrUtil.isNotBlank(content)) {
                            // 特殊处理system角色，确保它在第一条
                            if ("system".equals(role) && i > 0) {
                                // 如果system角色不在第一条，将其移到第一条
                                historyList.add(0, new RoleContent(role, content));
                            } else {
                                RoleContent roleContent = new RoleContent(role, content);
                                historyList.add(roleContent);
                            }
                        }
                    }
                }
            }
            
            // 构建请求参数（在添加用户消息到历史记录之前）
            JSONObject requestJson = buildRequestJson(text);
            
            // 将用户消息添加到历史记录中（除非是system角色，它应该已经在上下文中）
            boolean hasSystemRoleInContext = false;
            if (context != null) {
                for (int i = 0; i < context.size(); i++) {
                    JSONObject msg = context.get(i, JSONObject.class);
                    if ("system".equals(msg.getStr("role"))) {
                        hasSystemRoleInContext = true;
                        break;
                    }
                }
            }
            
            // 只有当上下文中没有system角色时，才添加用户消息到历史记录
            if (!hasSystemRoleInContext) {
                RoleContent userContent = new RoleContent("user", text);
                historyList.add(userContent);
            }

            webSocket.send(requestJson.toString());
        } catch (Exception e) {
            log.error("发送消息到讯飞AI异常", e);
        }
    }

    /**
     * 关闭连接
     */
    public void close() {
        if (webSocket != null) {
            webSocket.close(1000, "主动关闭连接");
            isConnected = false;
        }
    }

    /**
     * 构建请求JSON
     *
     * @param question 用户问题
     * @return 请求JSON
     */
    private JSONObject buildRequestJson(String question) {
        JSONObject jsonObject = JSONUtil.createObj();
        
        // header部分
        JSONObject header = JSONUtil.createObj();
        header.set("app_id", this.appId);
        header.set("uid", "ruoyi_ai_user");
        jsonObject.set("header", header);
        
        // parameter部分
        JSONObject parameter = JSONUtil.createObj();
        JSONObject chat = JSONUtil.createObj();
        
        // 根据host url判断domain
        String domain = "lite"; // 默认使用lite版本
        if (this.hostUrl != null) {
            if (this.hostUrl.contains("/v3.1/")) {
                domain = "generalv3"; // Spark Pro版本
            } else if (this.hostUrl.contains("/v3.5/")) {
                domain = "generalv3.5"; // Spark Max版本
            } else if (this.hostUrl.contains("/v1.1/")) {
                domain = "lite"; // Spark Lite版本
            }
        }
        
        chat.set("domain", domain);
        chat.set("temperature", 0.5);
        chat.set("max_tokens", 1024);
        parameter.set("chat", chat);
        jsonObject.set("parameter", parameter);
        
        // payload部分
        JSONObject payload = JSONUtil.createObj();
        JSONObject message = JSONUtil.createObj();
        
        // 构建消息内容，包含system角色和优化的上下文处理
        buildMessageContent(message, question);
        
        payload.set("message", message);
        jsonObject.set("payload", payload);
        
        return jsonObject;
    }

    /**
     * 构建消息内容，包含system角色和优化的上下文处理
     *
     * @param message 消息对象
     * @param question 用户问题
     */
    private void buildMessageContent(JSONObject message, String question) {
        List<JSONObject> textArray = new ArrayList<>();
        
        // 添加system角色（如果存在）
        if (!historyList.isEmpty() && "system".equals(historyList.get(0).role)) {
            RoleContent systemContent = historyList.get(0);
            JSONObject systemItem = JSONUtil.createObj();
            systemItem.set("role", systemContent.role);
            systemItem.set("content", systemContent.content);
            textArray.add(systemItem);
            
            // 从索引1开始添加历史记录，跳过system角色
            addHistoryItems(textArray, 1);
        } else {
            // 没有system角色，正常添加所有历史记录
            addHistoryItems(textArray, 0);
        }
        
        // 添加当前用户问题
        JSONObject userItem = JSONUtil.createObj();
        userItem.set("role", "user");
        userItem.set("content", question);
        textArray.add(userItem);
        
        message.set("text", textArray);
    }

    /**
     * 添加历史记录到消息内容
     *
     * @param textArray 文本数组
     * @param startIndex 开始索引
     */
    private void addHistoryItems(List<JSONObject> textArray, int startIndex) {
        synchronized (historyList) {
            // 计算实际可用的历史记录数量
            int availableHistory = historyList.size() - startIndex;
            // 限制最大历史记录数（避免超出上下文限制）
            int maxHistory = Math.min(availableHistory, 10); // 最多保留10轮对话
            
            // 添加历史记录
            for (int i = startIndex; i < historyList.size(); i++) {
                RoleContent roleContent = historyList.get(i);
                JSONObject item = JSONUtil.createObj();
                item.set("role", roleContent.role);
                item.set("content", roleContent.content);
                textArray.add(item);
            }
        }
    }

    /**
     * 获取鉴权URL
     *
     * @param hostUrl 讯飞API地址
     * @param apiKey API Key
     * @param apiSecret API Secret
     * @return 鉴权URL
     * @throws Exception 加密异常
     */
    public static String getAuthUrl(String hostUrl, String apiKey, String apiSecret) throws Exception {
        // 处理wss协议，OkHttp不支持直接解析wss协议
        String httpUrl = hostUrl;
        if (hostUrl.startsWith("wss://")) {
            httpUrl = "https://" + hostUrl.substring(6);
        } else if (hostUrl.startsWith("ws://")) {
            httpUrl = "http://" + hostUrl.substring(5);
        }

        okhttp3.HttpUrl url = okhttp3.HttpUrl.parse(httpUrl);
        // 添加null检查
        if (url == null) {
            throw new IllegalArgumentException("无效的hostUrl: " + hostUrl);
        }

        SimpleDateFormat format = new SimpleDateFormat("EEE, dd MMM yyyy HH:mm:ss z", Locale.US);
        format.setTimeZone(TimeZone.getTimeZone("GMT"));
        String date = format.format(new Date());

        String preStr = "host: " + url.host() + "\n" +
                "date: " + date + "\n" +
                "GET " + url.encodedPath() + " HTTP/1.1";

        Mac mac = Mac.getInstance("hmacsha256");
        SecretKeySpec spec = new SecretKeySpec(apiSecret.getBytes(StandardCharsets.UTF_8), "hmacsha256");
        mac.init(spec);

        byte[] hexDigits = mac.doFinal(preStr.getBytes(StandardCharsets.UTF_8));
        String sha = Base64.getEncoder().encodeToString(hexDigits);

        String authorization = String.format("api_key=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"",
                apiKey, "hmac-sha256", "host date request-line", sha);

        HttpUrl httpUrlObj = HttpUrl.parse("https://" + url.host() + url.encodedPath()).newBuilder()
                .addQueryParameter("authorization", Base64.getEncoder().encodeToString(authorization.getBytes(StandardCharsets.UTF_8)))
                .addQueryParameter("date", date)
                .addQueryParameter("host", url.host())
                .build();

        // 将https协议转换回wss协议
        String authUrl = httpUrlObj.toString();
        if (hostUrl.startsWith("wss://")) {
            authUrl = authUrl.replaceFirst("https://", "wss://");
        } else if (hostUrl.startsWith("ws://")) {
            authUrl = authUrl.replaceFirst("http://", "ws://");
        }

        return authUrl;
    }

    /**
     * 讯飞AI WebSocket监听器
     */
    private class XfWebSocketListener extends WebSocketListener {
        @Override
        public void onOpen(WebSocket webSocket, Response response) {
            log.info("与讯飞AI建立WebSocket连接成功");
        }

        @Override
        public void onMessage(WebSocket webSocket, String text) {
            try {
                // 解析讯飞AI返回的消息
                JSONObject jsonObject = JSONUtil.parseObj(text);
                if (jsonObject.containsKey("header")) {
                    JSONObject header = jsonObject.getJSONObject("header");
                    int code = header.getInt("code");
                    if (code != 0) {
                        log.error("讯飞AI返回错误码: {}, 错误信息: {}", code, header.getStr("message"));
                        webSocketSession.sendMessage(new org.springframework.web.socket.TextMessage("{\"code\": " + code + ", \"content\": \"" + header.getStr("message") + "\"}"));
                        XunfeiWebSocketClient.this.close();
                        return;
                    }
                }

                // 处理返回结果
                StringBuilder result = new StringBuilder();
                if (jsonObject.containsKey("payload") && jsonObject.getJSONObject("payload").containsKey("choices")) {
                    JSONObject payload = jsonObject.getJSONObject("payload");
                    JSONObject choices = payload.getJSONObject("choices");

                    if (choices.containsKey("text")) {
                        JSONArray textList = choices.getJSONArray("text");
                        for (int i = 0; i < textList.size(); i++) {
                            JSONObject item = textList.getJSONObject(i);
                            result.append(item.getStr("content"));
                        }
                    }
                }

                // 发送消息到前端WebSocket
                webSocketSession.sendMessage(new org.springframework.web.socket.TextMessage(result.toString()));

                // 将AI回复添加到历史记录中
                if (result.length() > 0) {
                    synchronized (historyList) {
                        RoleContent assistantContent = new RoleContent("assistant", result.toString());
                        historyList.add(assistantContent);
                    }
                }

                // 判断是否结束
                if (jsonObject.containsKey("header") && jsonObject.getJSONObject("header").containsKey("status")) {
                    int status = jsonObject.getJSONObject("header").getInt("status");
                    if (status == 2) {
                        // 讯飞AI返回结束标识
                        webSocketSession.sendMessage(new org.springframework.web.socket.TextMessage("[DONE]"));
                        XunfeiWebSocketClient.this.close();
                    }
                }
            } catch (Exception e) {
                log.error("处理讯飞AI返回消息异常", e);
            }
        }

        @Override
        public void onClosed(WebSocket webSocket, int code, String reason) {
            log.info("与讯飞AI的WebSocket连接已关闭, code: {}, reason: {}", code, reason);
            isConnected = false;
        }

        @Override
        public void onFailure(WebSocket webSocket, Throwable t, Response response) {
            log.error("讯飞AI WebSocket连接异常", t);
            isConnected = false;
        }
    }

    /**
     * 角色内容类
     */
    public static class RoleContent {
        public String role;
        public String content;

        public RoleContent() {
        }

        public RoleContent(String role, String content) {
            this.role = role;
            this.content = content;
        }
    }
}
